/***************************************************************************
*   Copyright (C) 2009 by Steven Waslander   *
*   stevenw@uwaterloo.ca   *
*                                                                         *
***************************************************************************/

#ifndef _MATRIX_H
#define _MATRIX_H

#include <iostream>
#include <fstream>
#include <cassert>
#include <cmath>
#include "Stats.h"

static const bool matrixDebug = true; 

template <class T>
class Matrix
{
  public:
    //! Constructor allocates necessary space on the heap
    Matrix(int rows = 3, int columns = 3);
    Matrix(const Matrix& src);
    //! Destructor deletes allocated memory
    ~Matrix();

    //! Matrix size accessors
    int NumRows(void) const;
    int NumCols(void) const;

    //! Nice for accessing matrix using A(1,2) notation
    inline T& operator() (int i, int j) // get and set 
    {
      if (matrixDebug) assert(i>=0 && j>=0 && i<rows && j<cols);
      return data[i*cols+j];
    }

    inline T& operator() (int i, int j) const // get only for const correctness
    {
      if (matrixDebug) assert(i>=0 && j>=0 && i< rows && j < cols);
      return data[i*cols+j];
    }
   
    //! Math operators
    Matrix<T>& operator = (const Matrix<T>& src);
    Matrix<T>& operator += (const Matrix<T>& src);
    Matrix<T>& operator -= (const Matrix<T>& src);
    
    Matrix<T> operator + (const Matrix<T>& right) const;
    Matrix<T> operator - (const Matrix<T>& right) const;
    Matrix<T> operator * (const Matrix<T>& right) const;
    Matrix<T> operator * (const double right) const;
    Matrix<T> operator / (const double right) const;

    //! Display operator
    template <class S>
    friend std::ostream& operator<< (std::ostream& out,const Matrix<S>& src);
    template <class S>
    friend std::istream& operator>> (std::istream& in, Matrix<S>& dest);
    template <class S>
    friend std::ifstream& operator>> (std::ifstream& in, Matrix<S>& dest);

    // Matrix functions
    //! Set entire matrix to a single value
     void SetVal(T val);
     //! Resize allocated space
     void Resize(int newrows,int newcols);
     void Resize(const Matrix<T>& target) ;
     //! Add rows, but preserve data
     void AddRows(int newrows);
     //! Add cols, but preserve data
     void AddCols(int newcols);
     //! Expand a matrix but preserve data
     void Grow(int newrows, int newcols);
     //! Transpose matrix in place
     Matrix<T> Transpose(void);
     //! Transpose matrix to a new matrix
     void Transpose(Matrix<T>& out) const;
     //! Calculate trace of matrix
     T Trace(void);
     //! Transforms matrix into identity of current size
     void eye();
     //! Transforms matrix into identity of size n
     void eye(int n);
  
    /*!Inserts a smaller source matrix into a larger target matrix at position (i,j). In
     * order to make sense, the bottom-right element of the src matrix cannot extend past
     * the bounds of the target matrix, when inserted at (i,j).
     * \param src The source matrix holding data that needs to be copied
     * \param i The row of the target matrix that src(0,0) should be copied to
     * \param j The col of the target matrix that src(0,0) should be copied to  
     */
    inline void Insert(const Matrix<T>& src, int i, int j) ;
  
    //! Reshapes matrix in the same sense as the Matlab reshape command
    static Matrix<T> Reshape(Matrix<T> src, int i, int j) ;
    Matrix<T>& Rand(T lower, T upper); //!< Make matrix elements uniform random
    Matrix<T>& Randn(T mean, T stdev); //!< Make matrix elements Gaussian
    void FullQR(Matrix<double>& Q1, Matrix<double>& Q2, Matrix<double>& R) ;       // Full QR Factorization
    Matrix<T> Null() ;
    
    // Square matrix operations
    T Det(void) const;
    int Inverse(void);
    int LuDecomp(Matrix<int>& indx, int* d); 
    void LuSolve(const Matrix<int>& indx, Matrix<T>& b);
    void pseudoInverse(Matrix<T>& out);

  protected:
    int rows, cols; //!< row, column size
    T* data; //!< Pointer to matrix data
};

template <class T>
T Norm(Matrix<T> x)
{
  T sum(0);
  if (x.NumCols()==1)
    for(int i=0; i<x.NumRows(); i++)
      sum += x(i,0)*x(i,0);
  else if (x.NumRows()==1)
    for(int i=0; i<x.NumRows(); i++)
      sum += x(i,0)*x(i,0);
  else
    assert(false); // Must have a single row or column
  return sqrt(sum);
}

//! Initializes without setting data values to anything, to ease allocation overhead
template <class T>
Matrix<T>::Matrix(int _rows, int _cols) :
  rows(_rows), cols(_cols)
{
  data = new T[rows*cols];
}

//! Copy constructor
template <class T>
Matrix<T>::Matrix(const Matrix<T>& src) : rows(src.rows), cols(src.cols)
{
  data = new T[rows*cols];

  for (int i=0;i<rows;i++) {
    for (int j=0;j<cols;j++) {
      (*this)(i,j) = src(i,j);
    }
  }
}

//! Wipes out allocated space on destruction
template <class T>
Matrix<T>::~Matrix()
{
  delete[] data;
}

template <class T>
void Matrix<T>::Insert(const Matrix<T>& src, int startRow, int startCol)
{
  assert(src.NumRows() + startRow <= NumRows()) ;                  // Verify that rows don't overflow
  assert(src.NumCols() + startCol <= NumCols()) ;                  // Verify that cols don't overflow
  for (int i = 0; i < src.NumRows(); i++)
    for (int j = 0; j < src.NumCols(); j++)
      (*this)(i+startRow, j+startCol) = src(i,j) ;
}

//! Sets every value in a matrix to the input
template <class T>
void Matrix<T>::SetVal(T val)
{
  for (int i=0; i<rows; i++){
    for (int j=0; j<cols; j++) {
      (*this)(i,j) = val;
    }
  }
}

//! Convert to identity matrix
template <class T>
void Matrix<T>::eye()
{
  (*this).SetVal(0);
  int size = rows>cols?cols:rows;
  for (int i=0; i<size; i++){
    (*this)(i,i) = 1;
  }
}

template <class T>
void Matrix<T>::eye(int n)
{
  (*this).Resize(n,n);
  (*this).eye();
}


template <class T>
int Matrix<T>::NumRows(void) const
{
  return rows;
}

template <class T>
int Matrix<T>::NumCols(void) const
{
  return cols;
}

template <class T>
Matrix<T> Matrix<T>::Reshape(Matrix<T> src, int numRows, int numCols)
{
  assert( numRows*numCols == src.NumRows()*src.NumCols() ) ;

  Matrix<T> dest(numRows,numCols) ;

  for(int elemNum=0; elemNum < numRows*numCols; elemNum++)
  {
    int srcCol = elemNum / src.NumRows() ;
    int srcRow = elemNum - srcCol*src.NumRows() ;
    
    int destCol = elemNum / dest.NumRows() ;
    int destRow = elemNum - destCol*dest.NumRows() ;
    dest(destRow, destCol) = src(srcRow, srcCol) ;
  }
  return dest ;
}

template <class T>
Matrix<T>& Matrix<T>::operator = (const Matrix<T>& src)
{
  if (this != &src) {
    assert(this->rows >=src.rows && this->cols >=src.cols);
    for (int i=0;i<src.rows;i++) {
      for (int j=0;j<src.cols; j++) {
        (*this)(i,j) = src(i,j);
      }
    }
  }
  return *this;
}

template <class T>
Matrix<T>& Matrix<T>::operator += (const Matrix<T>& src)
{
  assert(this->rows >=src.rows && this->cols >=src.cols);
  for (int i=0;i<src.rows;i++) {
    for (int j=0;j<src.cols;j++) {
      (*this)(i,j) += src(i,j);
    }
  }
  return *this;
}

template <class T>
Matrix<T>& Matrix<T>::operator -= (const Matrix<T>& src)
{
  assert(this->rows >=src.rows && this->cols >=src.cols);
  for (int i=0;i<src.rows;i++) {
    for (int j=0;j<src.cols;j++) {
      (*this)(i,j) -= src(i,j);
    }
  }
  return *this;
  
}

template <class T>
Matrix<T> Matrix<T>::operator + (const Matrix<T>& right) const
{
  Matrix<T> result(*this);
  result += right;
  return result;
}

template <class T>
Matrix<T> Matrix<T>::operator - (const Matrix<T>& right) const
{
  Matrix<T> result(*this);
  result -= right;
  return result;
}

template <class T>
Matrix<T> Matrix<T>::operator * (const Matrix<T>& right) const
{
  assert(this->cols == right.rows);
  Matrix<T> result(this->rows, right.cols);
  result.SetVal(0.0);
  for (int i=0;i<this->rows;i++) {
    for (int j=0;j<right.cols;j++) {
      for (int k=0;k<this->cols;k++) {
        result(i,j) += (*this)(i,k)*right(k,j);
      }
    }
  }
  return result;
}

template <class T>
Matrix<T> Matrix<T>::operator * (const double right) const
{
  Matrix<T> result(*this);
  for (int i=0; i < NumRows(); i++)
    for (int j=0; j < NumCols(); j++)
      result(i,j) = (*this)(i,j) * right ;
  return result ;
}

template <class T>
Matrix<T> Matrix<T>::operator / (const double right) const
{
  Matrix<T> result(*this);
  for (int i=0; i < NumRows(); i++)
    for (int j=0; j < NumCols(); j++)
      result(i,j) = (*this)(i,j) / right ;
  return result ;
}

//! Resizes the matrix to match the target matrix
template <class T>
void Matrix<T>::Resize(const Matrix<T>& target)
{
  Resize(target.NumRows(), target.NumCols()) ;
}


//! Resizes the matrix
template <class T>
void Matrix<T>::Resize(int newrows, int newcols)
{
  delete[] data;
  rows = newrows;
  cols = newcols;
  data = new T[rows*cols];
}

template <class T>
void Matrix<T>::AddRows(int newrows)
{
  Grow(newrows,0);
}

template <class T>
void Matrix<T>::AddCols(int newcols)
{
  Grow(0,newcols);
}

template <class T>
void Matrix<T>::Grow(int newrows,int newcols)
{
  Matrix<T> tmp(rows+newrows,cols+newcols);
  tmp.SetVal(0.0);  
  tmp.Insert((*this),0,0);
  
  delete[] data;
  rows = rows + newrows;
  cols = cols + newcols;
  data = new T[rows*cols];
  (*this) = tmp;
}



//! Transposes the matrix
template <class T>
Matrix<T> Matrix<T>::Transpose(void)
{
  Matrix<T> out(cols,rows);
  Transpose(out);
  return out;
}

//! Transposes the matrix, stores it to out, without allocating new space
template <class T>
void Matrix<T>::Transpose(Matrix<T>& out) const
{
  for (int i=0; i<rows; i++)
    for (int j=0; j<cols; j++)
      out(j,i) = (*this)(i,j);
}

//! Computes the trace of the matrix
template <class T>
T Matrix<T>::Trace(void)
{
  T sum = 0;
  for (int i=0; i<std::min(rows,cols); i++)
    sum += (*this)(i,i);
  return sum;
}

//! Fills the matrix with uniform random numbers between lower and upper
template <class T>
Matrix<T>& Matrix<T>::Rand(T lower, T upper)
{
  for (int i=0; i<(this->rows*this->cols); i++)
      this->data[i] = uniformRand(lower,upper);
  return *this;
}

//! Fills the matrix with Gaussian random numbers with mean and stdev
template <class T>
Matrix<T>& Matrix<T>::Randn(T mean, T stdev)
{
  for (int i=0; i<(this->rows*this->cols); i++)
      this->data[i] = normalRand(mean,stdev);
  return *this;
}

/** Return the determinant of this matrix.
 * This uses the LU decomposition.
 */
template <class T>
T Matrix<T>::Det(void) const
{
  Matrix<T> tmp(*this);
  Matrix<int> indx(rows,1);
  int d;
  T result;
  tmp.LuDecomp(indx,&d);
  result = T(d);
  for (int i=0; i<rows; i++)
    result *= tmp(i,i);
  return result;
}

/** Return the inverse of this matrix
 * This uses the LU decomposition
 */
template <class T> 
int Matrix<T>::Inverse(void)
{
  Matrix<T> tmp(*this);
  Matrix<int> indx(rows,1);
  Matrix<T> col(rows,1); 
  int d;

  if (tmp.LuDecomp(indx,&d)) {
    return 1;
  }
  for (int j=0;j<cols;j++) {
    col.SetVal(T(0));
    col(j,0) = T(1);
    tmp.LuSolve(indx,col);
    // Copy the result to the column
    for (int i=0; i<rows; i++)
      (*this)(i,j) = col(i,0);
  }
  return 0;
}


/***
 * Finds the Nullspace of the matrix using QR decomposition. Only works if A is square or fat.
 ***/
template <class T>
Matrix<T> Matrix<T>::Null()
{
  Matrix<double> Q1 ;
  Matrix<double> Q2 ;
  Matrix<double> R ;
  
  Matrix<double> At(this->Transpose()) ;
  At.FullQR(Q1,Q2,R) ;
  
  return Q2 ;
}


/***
 * Performs the full QR factorization on the current matrix.  Matrix must be square
 * or skinny for this to work
 * \todo Doesn't compute R yet.
 ***/
template <class T> 
void Matrix<T>::FullQR(Matrix<double>& Q1, Matrix<double>& Q2, Matrix<double>& R)
{
  //assert(NumRows() >= NumCols()) ;

  Matrix<double> A(NumRows(), NumRows()+NumCols()) ;
  Matrix<double> Eye(NumRows(), NumRows()) ;
  
  // Build identity matrix
  Eye.SetVal(0) ;
  for (int i=0; i < NumRows(); i++)
  {
    Eye(i,i) = 1 ;
  }
  A.Insert(*this,0,0) ;
  A.Insert(Eye, 0, NumCols()) ;
    
  // Perform the gram-schmidt decomp
  Matrix<double> Q(NumRows(), NumRows()) ;
  Q.SetVal(0) ;
  int Q2start = 0 ;
  int r = 0 ;       // Keeps track of which col in Q we are on
  Matrix<double> a(A.NumRows(), 1) ;    // Allocate memory to store a column of A
  Matrix<double> q(Q.NumRows(), 1) ;    // Allocate memory to store a column of A
  for (int i=0; i < A.NumCols(); i++)
  {
    // Get current column of A
    for (int k=0; k < A.NumRows(); k++)
      a(k,0) = A(k,i) ;
    // Subract off previous cols of Q
    for (int k=0; k < r; k++)
    {
      // Get current column of Q
      for (int j=0; j < Q.NumRows(); j++)
        q(j,0) = Q(j,k) ;
      Matrix<double> qa((q.Transpose()*a)) ;
      a = a - q * qa ;
    }
    
    if ((a.Transpose()*a)(0,0) > .0000001)
    {
      q = a * (1/sqrt( (a.Transpose() * a)(0,0) )) ;
      Q.Insert(q,0,r) ;
      r++ ;
    }
    if (i == NumCols()-1)
    {
      Q2start = r ;         // Need to mark when we've switched over to the columns of A that we appended
    }
  }
  
  Q1.Resize(NumRows(), Q2start) ;
  Q1.SetVal(0) ;
  for (int j=0; j<Q1.NumCols(); j++)
    for (int i=0; i<NumRows(); i++)
      Q1(i,j) = Q(i,j) ;
  Q2.Resize(NumRows(), Q.NumCols()-Q2start) ;
  Q2.SetVal(0) ;
  for (int j=Q2start; j<Q.NumCols(); j++)
    for (int i=0; i<NumRows(); i++)
      Q2(i,j-Q2start) = Q(i,j) ;
  
  R.SetVal(0) ;               //! \todo Need to figure out how to compute R for QR decomp
}

// Decompositions and solvers from Numerical Recipes in C 2nd Ed

/* 
 * @brief LU Decomposition in place
 * @return indx output vector that records the row permutation effected by 
 * the partial pivoting
 * @return d (optional) is +/- 1 depending on if there were an even (+) or 
 * odd (-) number of row-interchanges (used to compute determinant).
 * @returns 1 on failure (matrix singular), 0 on success
 */
template <class T> 
int Matrix<T>::LuDecomp(Matrix<int>& indx, int* d) 
{
  int i,imax=0,j,k; 
  T big,dum,sum,temp; 
  Matrix<T> vv(rows,1); // stores the implicit scaling of each row

  if (d) *d=1; // no row interchanges yet
  for (i=0;i<rows;i++) { // loop over rows to get the implicit scaling info 
    big=0.0; // initialize to smallest possible value 
    for (j=0;j<cols;j++) if ((temp=fabs((*this)(i,j))) > big) big=temp; 
    if (big == 0.0) return 1; // return error (matrix is singular)
    vv(i,0)=T(1)/big; // save the scaling
  }
  
  for (j=0;j<cols;j++) { // Loop over the columns for Crout's method 
    for (i=0;i<j;i++) { 
      sum=(*this)(i,j); 
      for (k=0;k<i;k++) sum -= (*this)(i,k)*(*this)(k,j); 
      (*this)(i,j)=sum; 
    } 
   
    big=0.0; // initialize for the search for the largest pivot element
    for (i=j;i<rows;i++) { 
      sum=(*this)(i,j); 
      for (k=0;k<j;k++) sum -= (*this)(i,k)*(*this)(k,j); 
      (*this)(i,j)=sum; 
      if ((dum=vv(i,0)*fabs(sum)) >= big) { 
        big=dum; 
        imax=i; //!todo check if imax is set here by initializing it to -1 first
      } 
    } 

    if (j != imax) { // check if row interchange is needed 
      for (k=0;k<cols;k++) {
        dum=(*this)(imax,k); 
        (*this)(imax,k)=(*this)(j,k); 
        (*this)(j,k)=dum; 
      }
      if (d) *d = -(*d); // change sign to indicate row swap occurred 
      vv(imax,0)=vv(j,0); 
    } 

    indx(j,0)=imax;
    if ((*this)(j,j) == 0.0) (*this)(j,j) = T(1e-20);
    if (j != rows-1) { 
      dum=1.0/((*this)(j,j)); 
      for (i=j+1;i<rows;i++) (*this)(i,j) *= dum; 
    } 
  }
  return 0;
}

/** Perform a LU Decomposition backsolve with given RHS.
 * Precondition: luDecomp(indx,d) has already been called
 * @param indx the row-swap vector returned from luDecomp()
 * @param b the RHS vector in the equation A*x = b
 * @return b is also the solution, overwriting the passed RHS
 */
template <class T> 
void Matrix<T>::LuSolve(const Matrix<int>& indx, Matrix<T>& b)
{
  int i,ii=-1,ip,j; 
  T sum;
  for (i=0; i<rows; i++) { 
    ip = indx(i,0);
    sum = b(ip,0);
    b(ip,0)=b(i,0); 
    if (ii >= 0)
      for (j=ii;j<i;j++)
        sum -= (*this)(i,j)*b(j,0); 
    else if (sum)
      ii=i;
    b(i,0) = sum; 
  } 
  for (i=(rows-1); i>=0; i--) { 
    sum = b(i,0); 
    for (j=i+1; j<rows; j++)
      sum -= (*this)(i,j)*b(j,0); 
    b(i,0) = sum/(*this)(i,i); 
  } 
}


template <class T> 
void Matrix<T>::pseudoInverse(Matrix<T>& out)
{  
  int r = (*this).NumRows();
  int c = (*this).NumCols();
  
  //! this transposed
  Matrix<T> thisT(c,r);
  (*this).Transpose(thisT);
    
  //! Check if least squares or least norm is needed
  if (r >= c) {
    //! this^T * this
    Matrix<T> A(c,c);
    A = thisT * (*this);
      
    //! Inverse of (this^T * this)
    Matrix<T> invA(A);
    invA.Inverse();
      
    //! Pseudo inverse
    out.Resize(c,r);
    out = invA * thisT;
      
  } else {
    //! this transposed * this
    Matrix<T> A(r,r);
    A = (*this) * thisT;
      
    //! Inverse of (T^T * T)
    Matrix<T> invA(A);
    invA.Inverse();
      
    //! Pseudo inverse
    out.Resize(c,r);
    out = thisT * invA;
  }
}  
    
template <class T>
std::ostream& operator<<(std::ostream& out, const Matrix<T>& src)
{
  out << "[ ";
  for (int i=0;i<src.rows;i++) {
    if (i>0)
      out << "  ";
    for (int j=0;j<src.cols;j++) {
      out.width(5);
      out << src(i,j);
      if (j<src.cols-1)
          out << ", ";
    }
    if (i<src.rows-1)
      out << ";" << std::endl;
  }
  out << "];" << std::endl;
  out << std::endl;
  return out;
}

template<class T>
std::istream& operator>>(std::istream& in, Matrix<T>& dest)
{
  char tmp;
  for (int i=0;i<dest.rows;i++) {
    in >> tmp;
    for (int j=0;j<dest.cols;j++) {
      in >> dest(i,j);
    }
    in >> tmp;
  }
  return in;
}

template<class T>
std::ifstream& operator>>(std::ifstream& in, Matrix<T>& dest)
{
  char tmp;
  //char line[256];
  //in.getline(line,256);
  in >> tmp;
  for (int i=0;i<dest.rows;i++) {
    for (int j=0;j<dest.cols;j++) {
      in >> dest(i,j);
      in >> tmp;
    }
  }
  return in;
}

#endif // _MATRIX_H
